from fastapi import APIRouter, Query, UploadFile, File,Form
import threading
import time
import asyncio
from starlette.responses import FileResponse
from typing import List
import uuid
import os
import embeddingDB
import kdbUtil
import oss2
from embeddingDB import pgDB
from concurrent.futures import ThreadPoolExecutor
# 创建一个线程池执行器
from fastapi.concurrency import run_in_threadpool
import asyncio
executor = ThreadPoolExecutor()
router = APIRouter(
    prefix= "/transfer",
    tags= ["transferRouter"],
    responses= {404: {"description": "Not found path knowledge"}}
)

# 阿里云OSS配置信息
OSS_ACCESS_KEY_ID = 'LTAI5t6UBGfq1KZyS3y3hUnT'
OSS_ACCESS_KEY_SECRET = '2wNzTY7hUVFTRrZCp818CLreAhAhEz'
auth = oss2.Auth(OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET)
# 设置Endpoint和Region
endpoint = "https://oss-cn-beijing.aliyuncs.com"
region = "cn-beijing"
bucket = oss2.Bucket(auth, endpoint, "zsai-algorithm", region=region)

@router.get("/downloadFile/{kdbTitle}/{url}")
def download(kdbTitle:str,url:str):
    # 处理完毕文件以后，生成了文件路径
    print(kdbTitle,url)
    # filename = "你要下载的文件路径.xls"
    path = os.path.join(os.path.dirname(os.path.realpath(__file__)),'uploadFiles',kdbTitle,url)
    print('下载文件的path',path)
    return FileResponse(
            path,  # 这里的文件名是你要发送的文件名
            filename=url.split('/')[-1], # 这里的文件名是你要给用户展示的下载的文件名，比如我这里叫lol.exe
        )

    
@router.post('/uploadFiles')
async def uploadFiles(fileList: List[UploadFile]=File(...),id:int=Form(...)):
    print(len(fileList),fileList,id)
    rs={"code":200,"msg":"上传成功","data":[]}
    saveUrls = []
    tasks = []
    
    tasks = [asyncio.create_task(process_single_file(file, id)) for file in fileList]
    
    # 等待所有任务完成
    await asyncio.gather(*tasks)
    #     # await run_in_threadpool(lambda:extractKnowledgeFromFile(file))
    # task = asyncio.create_task(extractKnowledgeFromFile(fileList,id))
    # tasks.append(task)
    # 等待所有任务完成
    # await asyncio.gather(*tasks)
    rs["data"] = ""
    return rs

async def process_single_file(file, id):
    await extract_knowledge_from_file(file, id)
    
async def extract_knowledge_from_file(file, id):
    print("进入提取知识")
    
    res = await file.read()  # 读取文件内容
    
    filename = file.filename
    file_id = 'file' + uuid.uuid4().hex
    new_filename = f"{filename.split('.')[0]}_{file_id}.{filename.split('.')[-1]}"
    root_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'uploadFiles', str(id))
    if not os.path.exists(root_dir):
        os.makedirs(root_dir)
        
    url = os.path.join(root_dir, new_filename)

    # 使用 run_in_threadpool 包装同步文件保存操作
    await run_in_threadpool(save_file, res, url)
    
    # 使用 run_in_threadpool 包装同步 OSS 文件上传操作
    oss_object_name = await run_in_threadpool(upload_to_oss, url, id, file_id, filename)
    print("oss_object_name::",oss_object_name)
    docs = kdbUtil.splitDocument(url)
    docs_dict = {
        "documents": docs,
        "fileName": [new_filename for _ in docs],
        "ids": [id for _ in docs],
        "oss_path": [oss_object_name for _ in docs]
    }
    
    print("准备持久化11111")
    # 使用 run_in_threadpool 包装同步数据库持久化操作
    await run_in_threadpool(persist_embedding, docs_dict)

def save_file(content, path):
    with open(path, "wb") as f:
        f.write(content)

def upload_to_oss(url, id, file_id, filename):
    with open(url, "rb") as f:
        oss_object_name = f"knowledgeDB/{id}/{file_id}.{filename.split('.')[-1]}"
        bucket.put_object(oss_object_name, f)
        return oss_object_name

def persist_embedding(docs):
    print("准备持久化2222")
    db = pgDB()
    db.persistentEmbedding(docs)
    db.close()